import numpy as np
import pyopencl as cl
from mako.template import Template
from math import ceil,sqrt
import hilbert

#neighborlist index: 0-n:normal, -1:null (end of)
#in cl kernel, judge -1 is faster than checking j<n_neigh or if? (can use const loop-size and unrolling)
#bond,angle,exclude _lists are static, can be done once at the beginning by CPU

class Cell:
    DENSITY_FLUCTUATION = 2.0 #=max_natoms_per_cell/ave_natom_per_cell
    cl_files = {'cell':'cell.cl'}

    def __init__(self,cfg,level=4,max_natoms_per_cell=None,name='cell',context=None,flags=''):
        nx = ny = nz = 2**level #must be 2^n
        bx,by,bz = cfg.box_l[0:3]   #box size
        dx,dy,dz = bx/nx,bx/ny,bz/nz #cell(bin) size
        rdx,rdy,rdz = 1./nx,1./ny,1./nz #reduced cell(bin) size
        ncells = nx*ny*nz #total
        if max_natoms_per_cell == None:
            ave_natoms_per_cell = cfg.n_atoms/float(ncells)
            max_natoms_per_cell = int(ceil(ave_natoms_per_cell*Cell.DENSITY_FLUCTUATION))

        self.max_natoms_per_cell = max_natoms_per_cell
        self.atoms_in_cell = np.ones((ncells,max_natoms_per_cell),'int32')*(-1)
        self.n_atoms_in_cell = np.zeros((ncells,),'int32')
        self.nx,self.ny,self.nz = nx,ny,nz
        self.dx,self.dy,self.dz = dx,dy,dz
        self.ncells = ncells
        self.to_hilbert,self.from_hilbert = self.get_hilbert_curve()

        self.name = name
        cl_file = Cell.cl_files[name]
        cl_code = file(cl_file,'r').read()
        cl_code = Template(cl_code).render(natoms_per_cell=max_natoms_per_cell,
                nx=nx,ny=ny,nz=nz)
        self.prg = cl.Program(context,cl_code).build(options=flags)
        mf = cl.mem_flags
        self._atoms_in_cell = cl.Buffer(context,mf.READ_WRITE,
                size=self.atoms_in_cell.nbytes)
        self._n_atoms_in_cell = cl.Buffer(context,mf.READ_WRITE,
                size=self.n_atoms_in_cell.nbytes)
        self.overflow = np.zeros((1,),'int32')
        self._overflow = cl.Buffer(context,mf.READ_WRITE,size=4)

    def sort_on_host(self,cfg):
        r_t = np.ndarray((cfg.n_atoms,4),'float32')
        order = np.ndarray((cfg.n_atoms,),'int32') #
        n = 0
        for i in range(self.ncells):
            ic = self.from_hilbert[i] #cells in hilbert order
            for i in range(self.n_atoms_in_cell[ic]):
                i_old = self.atoms_in_cell[ic,i]
                order[n] = i_old
                r_t[n] = cfg.r_t[i_old]
                n += 1
        cfg.order = order
        cfg.r_t = r_t
    
    def device_to_host(self,que):
        cl.enqueue_copy(que,self.atoms_in_cell,self._atoms_in_cell)
        cl.enqueue_copy(que,self.n_atoms_in_cell,self._n_atoms_in_cell)
        cl.enqueue_copy(que,self.overflow,self._overflow)
        print 'overflow:',self.overflow

    def cal_cell(self,que,sglb,sloc,cfg):
        self.prg.init_cell(que,sglb,sloc,
                self._atoms_in_cell,
                self._n_atoms_in_cell)
        self.prg.cal_cell(que,sglb,sloc,
                np.int32(cfg.n_atoms),cfg.box_lo,cfg.box_hi,
                cfg._r_t,cfg._rr_c,
                self._atoms_in_cell,
                self._n_atoms_in_cell,
                self._overflow)

    def get_hilbert_curve(self): #nomral --> hilbert
        to_hilbert = np.ndarray((self.ncells,),'int32')
        from_hilbert = np.ndarray((self.ncells,),'int32')
        for i in range(self.nx):
            for j in range(self.ny):
                for k in range(self.nz):
                    d_normal = i*self.ny*self.nz + j*self.nz +k
                    d_hilbert = hilbert.Hilbert_to_int([i,j,k])
                    to_hilbert[d_normal] = d_hilbert
                    from_hilbert[d_hilbert] = d_normal
        return to_hilbert,from_hilbert

    def create_cell_neighbors_sphere(self,cfg,rcut_cell):
        #this is a slow process, however, it is only need to calculate once
        #rcut_cell = rcut + rskin
        r_toler = sqrt(self.dx*self.dx + self.dy*self.dy + self.dz*self.dz)
        r_toler *= 1.1 #for safe cutoff
        r2cut = (rcut_cell+r_toler)*(rcut_cell+r_toler)
        nx,ny,nz = self.nx,self.ny,self.nz
        bx,by,bz = cfg.box_l[0:3]
        cell_neighbors_of_cell = []
        n_cell_neighbors_of_cell = []
        hx = int(ceil(rcut_cell/self.dx)) #half serach width in x
        hy = int(ceil(rcut_cell/self.dy))
        hz = int(ceil(rcut_cell/self.dz))
        #print 'hx,hy,hz:',hx,hy,hz
        for i in range(self.ncells):
            cell_neighbors_of_cell.append([])
            i0,j0,k0 = self.d2ijk_cell(i)
            for i1 in range(i0-hx,i0+hx+1):
                for j1 in range(j0-hy,j0+hy+1):
                    for k1 in range(k0-hz,k0+hz+1):
                        if (i0,j0,k0) == (i1,j1,k1):
                            continue
                        dx = (i1-i0)*self.dx
                        dy = (j1-j0)*self.dy
                        dz = (k1-k0)*self.dz
                        dx -= round(dx/bx)*bx*cfg.box_p[0]
                        dy -= round(dy/by)*by*cfg.box_p[1]
                        dz -= round(dz/bz)*bz*cfg.box_p[2]
                        d2 = dx*dx + dy*dy + dz*dz
                        if d2>r2cut:
                            continue
                        i1 = (i1+nx)%nx
                        j1 = (j1+ny)%ny
                        k1 = (k1+nz)%nz
                        j = self.ijk2d_cell(i1,j1,k1)
                        cell_neighbors_of_cell[i].append(j)
                        #cell_neighbors_of_cell[i].append([i1,j1,k1])
        self.cell_neighbors_of_cell = cell_neighbors_of_cell 
        for c in cell_neighbors_of_cell:
            n_cell_neighbors_of_cell.append(len(c))
        self.n_cell_neighbors_of_cell = n_cell_neighbors_of_cell 
        self.max_n_cell_neighbors_per_cell = max(n_cell_neighbors_of_cell)

    def create_cell_neighbors_cube(self,cfg,rcut_cell):
        #this is a slow process, however, it is only need to calculate once
        #rcut_cell = rcut + rskin
        r2cut = rcut_cell*rcut_cell
        nx,ny,nz = self.nx,self.ny,self.nz
        cell_neighbors_of_cell = []
        n_cell_neighbors_of_cell = []
        hx = int(ceil(rcut_cell/self.dx)) #half serach width in x
        hy = int(ceil(rcut_cell/self.dy))
        hz = int(ceil(rcut_cell/self.dz))
        #print 'hx,hy,hz:',hx,hy,hz
        for i in range(self.ncells):
            cell_neighbors_of_cell.append([])
            i0,j0,k0 = self.d2ijk_cell(i)
            for i1 in range(i0-hx,i0+hx+1):
                for j1 in range(j0-hy,j0+hy+1):
                    for k1 in range(k0-hz,k0+hz+1):
                        if (i0,j0,k0) == (i1,j1,k1):
                            continue
                        i1 = (i1+nx)%nx
                        j1 = (j1+ny)%ny
                        k1 = (k1+nz)%nz
                        j = self.ijk2d_cell(i1,j1,k1)
                        cell_neighbors_of_cell[i].append(j)
                        #cell_neighbors_of_cell[i].append([i1,j1,k1])
        self.cell_neighbors_of_cell = cell_neighbors_of_cell 
        for c in cell_neighbors_of_cell:
            n_cell_neighbors_of_cell.append(len(c))
        self.n_cell_neighbors_of_cell = n_cell_neighbors_of_cell 

    def d2ijk_cell(self,d):
        nx,ny,nz = self.nx,self.ny,self.nz
        k = d % self.nz
        j = (d-k)/nz % ny
        i = ((d-k)/nz-j)/ny % nx
        return i,j,k

    def ijk2d_cell(self,i,j,k):
        nx,ny,nz = self.nx,self.ny,self.nz
        d = i*ny*nz + j*nz + k
        return d

